
#导入所需的模块
import numpy as np
import cv2
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import pickle

def abs_sobel_thresh(img, orient='x', thresh_min=20, thresh_max=100):
    '''
    获取图像，渐变方向和阈值的最小值与最大值
    :param img:
    :param orient:
    :param thresh_min:
    :param thresh_max:
    :return:
    '''
    #转换为灰度
    #cv2.cvtColor() 颜色空间转换函数
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    #使用OpenCV的Sobel()函数应用x或y渐变
    #取绝对值
    if orient == 'x':
        #利用Sobel方法可以进行sobel边缘检测
        #gary表示源图像，即进行边缘检测的图像
        #cv2.CV_64F表示64位浮点数即64 float
        #第三和第四个参数分别是对x和y方向的导数，即dx,dy
        #对于图像来说就是差分，这里1表示对x求偏导，0表示不对y求导
        #对x求导就是检测x方向上是否有边缘
        abs_sobel = np.absolute(cv2.Sobel(gray, cv2.CV_64F,1,0))
    if orient == 'y':
        abs_sobel = np.absolute(cv2.Sobel(gray, cv2.CV_64F,0,1))

    #重新缩放回8位整数
    scaled_sobel = np.uint8(255 * abs_sobel / np.max(abs_sobel))

    #创建一个副本并应用阈值
    binary_output = np.zeros_like(scaled_sobel)
    #这里我使用 >= and <= 阈值，但是< and >也可以
    binary_output[(scaled_sobel >= thresh_min) & (scaled_sobel <= thresh_max)] = 1
    #返回结果
    return binary_output

def mag_thresh(img, sobel_kernel=3, mag_thresh=(30,100)):
    #返回给定Sobel内核大小和阈值的梯度大小
    #转换为灰度
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    #取Sobel x和 y梯度
    sobelx = cv2.Sobel(gray, cv2.CV_64F,1,0,ksize=sobel_kernel)
    sobely = cv2.Sobel(gray, cv2.CV_64F,0,1,ksize=sobel_kernel)
    #计算梯度大小
    gradmag = np.sqrt(sobelx ** 2 + sobely ** 2)
    #重新缩放到8位
    scale_factor = np.max(gradmag)/ 255
    gradmag = (gradmag / scale_factor).astype(np.uint8)
    #创建二值图像达到阈值为1，否则为0
    binary_output = np.zeros_like(gradmag)
    binary_output[(gradmag >= mag_thresh[0])& (gradmag <= mag_thresh[1])] = 1

    #返回二值图像
    return binary_output

def dir_threshold(img, sobel_kernel=3, thresh=(0,np.pi / 2)):
    '''
    返回给定的sobel内核大小和阈值渐变的方向
    :param img:
    :param sobel_kernel:
    :param thresh:
    :return:
    '''
    #转换为灰度
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    #计算x和y梯度
    sobelx = cv2.Sobel(gray, cv2.CV_64F,1,0,ksize=sobel_kernel)
    sobely = cv2.Sobel(gray, cv2.CV_64F,0,1,ksize=sobel_kernel)

    #取梯度方向的绝对值
    #应用阈值，并创建一个二值图像结果
    absgraddir = np.arctan(np.arctan(np.absolute(sobely), np.absolute(sobelx)))
    binary_output = np.zeros_like(absgraddir)
    binary_output[(absgraddir >= thresh[0]) & (absgraddir <= thresh[1])] = 1

    #返回二值图像
    return binary_output

def hls_thresh(img, thresh=(100,255)):
    '''
    利用S通道将RGB转换为HLS,将阈值转换为二值图像
    :param img:
    :param thresh:
    :return:
    '''
    hls = cv2.cvtColor(img, cv2.COLOR_RGB2HLS)
    s_channel = hls[:,:,2]
    binary_output = np.zeros_like(s_channel)
    binary_output[(s_channel > thresh[0]) & (s_channel <= thresh[1])] = 1
    return binary_output

def combined_thresh(img):
    abs_bin = abs_sobel_thresh(img, orient='x', thresh_min=50, thresh_max=255)
    mag_bin = mag_thresh(img, sobel_kernel=3, mag_thresh=(50,255))
    dir_bin = dir_threshold(img, sobel_kernel=15, thresh=(0.7,1.3))
    hls_bin = hls_thresh(img, thresh=(170,255))

    combined = np.zeros_like(dir_bin)
    combined[(abs_bin ==  1 | ((mag_bin ==  1) & (dir_bin == 1))) | hls_bin == 1] = 1
    return combined, abs_bin, mag_bin, dir_bin, hls_bin  #调试

if __name__ == '__main__':
    img_file = 'test_images/test6.jpg'

    with open('calibrate_camera.p', 'rb') as f:
        save_dict = pickle.load(f)

    mtx = save_dict['mtx']
    dist = save_dict['dist']

    img = mpimg.imread(img_file)
    img = cv2.undistort(img, mtx, dist, None, mtx)

    combined, abs_bin, mag_bin, dir_bin, hls_bin = combined_thresh(img)

    plt.subplot(2,3,1)
    plt.imshow(abs_bin, cmap='gray', vmin=0, vmax=1)
    plt.subplot(2,3,2)
    plt.imshow(mag_bin, cmap='gray', vmin=0, vmax=1)
    plt.subplot(2,3,3)
    plt.imshow(dir_bin, cmap='gray', vmin=0, vmax=1)
    plt.subplot(2,3,4)
    plt.imshow(hls_bin, cmap='gray', vmin=0, vmax=1)
    plt.subplot(2,3,5)
    plt.imshow(img)
    plt.subplot(2,3,6)
    plt.imshow(combined, cmap='gray', vmin=0, vmax=1)

    plt.tight_layout()
    plt.show()






